# State vector module for FLEXINVERT.jl
# Ports mod_states.f90 functionality

module StateVector

using ..CoreTypes: Domain
using ..Settings: Config
using ..PriorCovariance: CovarianceMatrix
using ..Transformations

export InversionState, initialize_state, update_state!, get_physical_state, get_control_state

"""
    InversionState

Represents the state vector for the inversion system with both control and physical representations.
"""
mutable struct InversionState
    # Control variables (reduced space)
    phi::Vector{Float64}              # Current control state
    phi_prior::Vector{Float64}        # Prior control state (usually zeros)

    # Physical variables (full space)
    chi::Vector{Float64}              # Current physical state
    chi_prior::Vector{Float64}        # Prior physical state

    # State vector metadata
    n_control::Int                    # Number of control variables
    n_physical::Int                   # Number of physical variables
    n_boxes::Int                      # Number of spatial boxes
    n_time::Int                       # Number of time steps

    # Covariance matrix for transforms
    covariance::CovarianceMatrix

    # State history for optimization
    iteration::Int
    cost_history::Vector{Float64}
    gradient_norm_history::Vector{Float64}
end

"""
    initialize_state(cfg::Config, domain::Domain, covariance::CovarianceMatrix) -> InversionState

Initialize the state vector with prior values and covariance structure.

# Arguments
- `cfg`: Configuration settings
- `domain`: Domain information
- `covariance`: Prior covariance matrix with eigenvalue decomposition

# Returns
- Initialized InversionState
"""
function initialize_state(cfg::Config, domain::Domain, covariance::CovarianceMatrix)
    @info "Initializing inversion state"

    n_control = covariance.n_modes
    n_physical = covariance.nbox * covariance.nt
    n_boxes = domain.nbox
    n_time = covariance.nt

    @info "State dimensions: $n_control control variables, $n_physical physical variables"

    # Initialize control variables (usually start at zero = prior)
    phi = zeros(Float64, n_control)
    phi_prior = zeros(Float64, n_control)  # Prior is at the origin in control space

    # Initialize physical variables
    chi_prior = read_prior_state(cfg, domain, covariance)
    chi = copy(chi_prior)  # Start at prior

    # Initialize optimization history
    cost_history = Float64[]
    gradient_norm_history = Float64[]

    state = InversionState(
        phi, phi_prior,
        chi, chi_prior,
        n_control, n_physical, n_boxes, n_time,
        covariance,
        0, cost_history, gradient_norm_history
    )

    @info "State initialized with prior values"
    return state
end

"""
    read_prior_state(cfg::Config, domain::Domain, covariance::CovarianceMatrix) -> Vector{Float64}

Read prior state values from configuration or use defaults.

# Arguments
- `cfg`: Configuration settings
- `domain`: Domain information
- `covariance`: Prior covariance matrix (provides temporal dimension)

# Returns
- Vector of prior state values in physical space
"""
function read_prior_state(cfg::Config, domain::Domain, covariance::CovarianceMatrix)
    n_boxes = domain.nbox
    n_time = max(covariance.nt, 1)

    # Default prior flux scaling (deviation from zero)
    prior_scaling = hasfield(typeof(cfg), :prior_scaling) ? cfg.prior_scaling : 0.0

    # Create prior state vector
    prior_state = zeros(Float64, n_boxes * n_time)

    if prior_scaling != 0.0
        # Apply scaling proportional to grid cell area
        for t in 1:n_time
            for i in 1:n_boxes
                idx = (t - 1) * n_boxes + i
                area = domain.area_box[i]
                prior_state[idx] = prior_scaling * area
            end
        end
        @info "Applied prior scaling $prior_scaling to state vector"
    else
        @info "Using zero prior state vector"
    end

    return prior_state
end

"""
    update_state!(state::InversionState, dphi::Vector{Float64})

Update the state vector with control space increment.

# Arguments
- `state`: Current inversion state (modified in-place)
- `dphi`: Increment in control space
"""
function update_state!(state::InversionState, dphi::Vector{Float64})
    if length(dphi) != state.n_control
        error("Control increment length ($(length(dphi))) must match state dimension ($(state.n_control))")
    end

    # Update control variables
    state.phi .+= dphi

    # Transform to physical space
    Transformations.phi2chi!(state.chi, state.phi, state.covariance)

    # Add prior (since control space is centered at zero)
    state.chi .+= state.chi_prior

    @debug "State updated: control norm = $(norm(state.phi)), physical norm = $(norm(state.chi))"
end

"""
    get_physical_state(state::InversionState) -> Vector{Float64}

Get current state in physical space.

# Arguments
- `state`: Inversion state

# Returns
- Current state vector in physical space
"""
function get_physical_state(state::InversionState)
    return copy(state.chi)
end

"""
    get_control_state(state::InversionState) -> Vector{Float64}

Get current state in control space.

# Arguments
- `state`: Inversion state

# Returns
- Current state vector in control space
"""
function get_control_state(state::InversionState)
    return copy(state.phi)
end

"""
    reset_to_prior!(state::InversionState)

Reset state to prior values.

# Arguments
- `state`: Inversion state (modified in-place)
"""
function reset_to_prior!(state::InversionState)
    state.phi .= state.phi_prior
    state.chi .= state.chi_prior
    state.iteration = 0
    empty!(state.cost_history)
    empty!(state.gradient_norm_history)

    @info "State reset to prior"
end

"""
    set_from_physical!(state::InversionState, chi_new::Vector{Float64})

Set state from physical space values.

# Arguments
- `state`: Inversion state (modified in-place)
- `chi_new`: New physical state values
"""
function set_from_physical!(state::InversionState, chi_new::Vector{Float64})
    if length(chi_new) != state.n_physical
        error("Physical state length ($(length(chi_new))) must match state dimension ($(state.n_physical))")
    end

    # Set physical state
    state.chi .= chi_new

    # Remove prior and transform to control space
    chi_deviation = chi_new .- state.chi_prior
    Transformations.chi2phi!(state.phi, chi_deviation, state.covariance)

    @debug "State set from physical: control norm = $(norm(state.phi)), physical norm = $(norm(state.chi))"
end

"""
    add_to_optimization_history!(state::InversionState, cost::Float64, grad_norm::Float64)

Add optimization metrics to history.

# Arguments
- `state`: Inversion state (modified in-place)
- `cost`: Current cost function value
- `grad_norm`: Current gradient norm
"""
function add_to_optimization_history!(state::InversionState, cost::Float64, grad_norm::Float64)
    state.iteration += 1
    push!(state.cost_history, cost)
    push!(state.gradient_norm_history, grad_norm)

    @info "Iteration $(state.iteration): cost = $cost, |grad| = $grad_norm"
end

"""
    get_state_statistics(state::InversionState) -> NamedTuple

Get statistics about the current state.

# Arguments
- `state`: Inversion state

# Returns
- NamedTuple with state statistics
"""
function get_state_statistics(state::InversionState)
    # Control space statistics
    phi_norm = norm(state.phi)
    phi_max = maximum(abs.(state.phi))

    # Physical space statistics
    chi_norm = norm(state.chi)
    chi_max = maximum(abs.(state.chi))
    chi_mean = mean(state.chi)

    # Deviation from prior
    chi_deviation = state.chi .- state.chi_prior
    deviation_norm = norm(chi_deviation)

    return (
        iteration = state.iteration,
        phi_norm = phi_norm,
        phi_max = phi_max,
        chi_norm = chi_norm,
        chi_max = chi_max,
        chi_mean = chi_mean,
        deviation_from_prior = deviation_norm,
        n_control = state.n_control,
        n_physical = state.n_physical
    )
end

"""
    validate_state_consistency(state::InversionState; rtol::Float64=1e-10) -> Bool

Validate that control and physical representations are consistent.

# Arguments
- `state`: Inversion state
- `rtol`: Relative tolerance for validation

# Returns
- `true` if state is consistent, `false` otherwise
"""
function validate_state_consistency(state::InversionState; rtol::Float64=1e-10)
    # Transform control to physical and compare
    chi_from_phi = Transformations.phi2chi(state.phi, state.covariance) .+ state.chi_prior

    error_norm = norm(chi_from_phi .- state.chi)
    chi_norm = norm(state.chi)

    relative_error = chi_norm > 0 ? error_norm / chi_norm : error_norm

    if relative_error > rtol
        @warn "State inconsistency detected: relative error = $relative_error > $rtol"
        return false
    end

    return true
end

# Copy methods for MCMC integration
"""
    Base.copy(state::InversionState) -> InversionState

Create a copy of the InversionState for MCMC proposals.
"""
function Base.copy(state::InversionState)
    return InversionState(
        copy(state.phi),
        copy(state.phi_prior),
        copy(state.chi),
        copy(state.chi_prior),
        state.n_control,
        state.n_physical,
        state.n_boxes,
        state.n_time,
        state.covariance,  # Covariance is immutable, no need to copy
        state.iteration,
        copy(state.cost_history),
        copy(state.gradient_norm_history)
    )
end

"""
    Base.copyto!(dest::InversionState, src::InversionState) -> InversionState

Copy data from source to destination InversionState.
"""
function Base.copyto!(dest::InversionState, src::InversionState)
    # Ensure dimensions match
    if dest.n_control != src.n_control || dest.n_physical != src.n_physical
        error("Cannot copy between InversionStates with different dimensions")
    end

    # Copy arrays
    copyto!(dest.phi, src.phi)
    copyto!(dest.phi_prior, src.phi_prior)
    copyto!(dest.chi, src.chi)
    copyto!(dest.chi_prior, src.chi_prior)

    # Copy scalar values
    dest.iteration = src.iteration

    # Copy histories
    empty!(dest.cost_history)
    append!(dest.cost_history, src.cost_history)
    empty!(dest.gradient_norm_history)
    append!(dest.gradient_norm_history, src.gradient_norm_history)

    return dest
end

end # module StateVector
